import base64
import traceback

from ..common import VolcengineClient
from ..util import convert_bytes_to_image, convert_image_to_bytes

versions = ["通用2.0", "通用1.4"]


class KeepFaceSwap:
    RETURN_TYPES = ("IMAGE",)
    # RETURN_NAMES = ("image_output_name",)

    FUNCTION = "image_creator"

    # OUTPUT_NODE = False

    CATEGORY = "KeepFaceSwap"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "info": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "{}",
                    "lazy": True
                }),
                "image": ("IMAGE",),
                "version": (versions, {"default": versions[0]}),
                "prompt": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "",
                    "lazy": True
                }),
                "scale": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100, "step": 0.05}),
                "cfg_rescale": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 100, "step": 0.05}),
                "ref_ip_weight": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 100, "step": 0.01}),
                "ref_id_weight": ("FLOAT", {"default": 0.36, "min": 0.0, "max": 100, "step": 0.01}),
                "steps": ("INT", {"default": 90, "min": 0, "max": 200, "step": 1}),
                "width": ("INT", {"default": 512, "min": 0, "max": 10000, "step": 1}),
                "height": ("INT", {"default": 512, "min": 0, "max": 10000, "step": 1}),
            },
        }

    @staticmethod
    def image_creator(info, image, version, prompt, scale, ref_ip_weight, ref_id_weight, steps, width, height,
                      cfg_rescale):
        try:
            client = VolcengineClient(volcengine_info=info)
            if isinstance(image, list):
                images = []
                for img in image:
                    image_bytes = convert_image_to_bytes(img)
                    img_base64 = base64.b64encode(image_bytes).decode("utf-8")
                    if version == "通用1.4":
                        operate_result = client.keep_face_14(image_base64s=[img_base64], prompt=prompt, scale=scale,
                                                             steps=steps, width=width, height=height,
                                                             ref_id_weight=ref_id_weight,
                                                             ref_ip_weight=ref_ip_weight)
                    else:
                        operate_result = client.keep_face_20(image_base64s=[img_base64], prompt=prompt, scale=scale,
                                                             steps=steps, width=width, height=height,
                                                             cfg_rescale=cfg_rescale, ref_id_weight=ref_id_weight,
                                                             ref_ip_weight=ref_ip_weight)
                    if operate_result.code == 1:
                        pics = operate_result.data
                        image_bytes = pics[0]
                        pic = convert_bytes_to_image(image_bytes)
                        images.append(pic)
                image = images
            else:
                image_bytes = convert_image_to_bytes(image)
                image_base64 = base64.b64encode(image_bytes).decode("utf-8")
                if version == "通用1.4":
                    operate_result = client.keep_face_14(image_base64s=[image_base64], prompt=prompt, scale=scale,
                                                         steps=steps, width=width, height=height,
                                                         ref_id_weight=ref_id_weight,
                                                         ref_ip_weight=ref_ip_weight)
                else:
                    operate_result = client.keep_face_20(image_base64s=[image_base64], prompt=prompt, scale=scale,
                                                         steps=steps, width=width, height=height,
                                                         cfg_rescale=cfg_rescale, ref_id_weight=ref_id_weight,
                                                         ref_ip_weight=ref_ip_weight)
                if operate_result.code == 1:
                    pics = operate_result.data
                    image_bytes = pics[0]
                    image = convert_bytes_to_image(image_bytes)
        except Exception as e:
            print(e, traceback.format_exc())
        return (image,)
